package bin_tree.heap.leetcode;

import java.util.Comparator;
import java.util.PriorityQueue;
import java.util.Queue;

/**
 * 面试题17.14 找到最小的K个元素
 *
 * @author shkstart
 * @create 2022-10-29-21:21
 */
public class TopK {

    public int[] smallestK(int[] arr, int k) {

        int[] ret = new int[k];
//        边界条件
        if (arr.length == 0 || k == 0) {
            return ret;
        }
//        构建最大堆
//        JDK 中内置的优先级队列是最小堆，我们需要将其转为最大堆
//        此时我们用到比较器
//        从元素越小优先级越高的最小堆变为 元素越大的优先级越高的最大堆
        Queue<Integer> queue = new PriorityQueue<>(new Comparator<Integer>() {
            @Override
            public int compare(Integer o1, Integer o2) {
                return o2 - o1;
            }
        });

        for (int value : arr) {
            if (queue.size() < k) {
                queue.offer(value);
            } else {
                if (value < queue.peek()) {
                    queue.poll();
                    queue.offer(value);
                }
            }
        }

//        扫描完成后，queue中存储了前K小个元素
        for (int i = 0; i < k; i++) {
            ret[i] = queue.poll();

        }
        return ret;

    }
}
